#!/usr/bin/env python
# coding: utf-8

# In[4]:


import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import warnings
warnings.filterwarnings('ignore')

class QuantSystem:
    def __init__(self, stock_list, start_date, end_date, token):
        # 初始化Tushare接口
        ts.set_token(token)
        self.pro = ts.pro_api()
        
        self.stock_list = stock_list
        self.start_date = start_date
        self.end_date = end_date
        self.strategies = {
            'price_volume': self.price_volume_strategy,
            'financial': self.financial_strategy,
            'ml': self.ml_strategy
        }
        self.results = {}

    def fetch_data(self, ticker):
        """通过Tushare获取股票数据和财务指标"""
        try:
            # 获取日线行情数据
            df_daily = ts.pro_bar(ts_code=ticker, adj='qfq', 
                                 start_date=self.start_date, 
                                 end_date=self.end_date)
            if df_daily is None or df_daily.empty:
                print(f"无法获取{ticker}的日线行情数据")
                return None
                
            # Tushare返回的数据是倒序的，需要正序排列
            df_daily = df_daily.sort_values('trade_date')
            df_daily['trade_date'] = pd.to_datetime(df_daily['trade_date'])
            df_daily.set_index('trade_date', inplace=True)
            
            # 获取财务指标数据
            # 由于Tushare财务数据获取限制，这里只获取部分关键指标
            df_financial = self.get_financial_indicators(ticker)
            if df_financial is not None and not df_financial.empty:
                # 合并行情数据和财务数据
                df = pd.merge_asof(df_daily, df_financial, 
                                  left_index=True, right_index=True, 
                                  tolerance=pd.Timedelta('90D'), direction='backward')
                return df
            return df_daily
            
        except Exception as e:
            print(f"获取{ticker}数据时出错: {e}")
            return None

    def get_financial_indicators(self, ticker):
        """获取关键财务指标"""
        try:
            # 获取基本财务指标
            df_indicators = self.pro.fina_indicator(ts_code=ticker, 
                                                  start_date=self.start_date, 
                                                  end_date=self.end_date)
            if df_indicators is None or df_indicators.empty:
                return None
                
            df_indicators['end_date'] = pd.to_datetime(df_indicators['end_date'])
            df_indicators.set_index('end_date', inplace=True)
            df_indicators.sort_index(inplace=True)
            
            # 选择需要的财务指标
            selected_columns = ['roe', 'grossprofit_margin', 'netprofit_margin', 
                               'assets_turn', 'debt_to_assets', 'eps']
            df_selected = df_indicators[selected_columns]
            
            # 重命名列名以便于理解
            df_selected.rename(columns={
                'roe': 'ROE',
                'grossprofit_margin': 'Gross_Profit_Margin',
                'netprofit_margin': 'Net_Profit_Margin',
                'assets_turn': 'Assets_Turnover',
                'debt_to_assets': 'Debt_To_Assets',
                'eps': 'EPS'
            }, inplace=True)
            
            return df_selected
            
        except Exception as e:
            print(f"获取{ticker}财务指标时出错: {e}")
            return None

    def calculate_technical_indicators(self, df):
        """计算技术指标"""
        # 确保数据包含所需列
        required_columns = ['close', 'high', 'low', 'vol']
        if not all(col in df.columns for col in required_columns):
            missing = [col for col in required_columns if col not in df.columns]
            print(f"缺少技术指标计算所需的列: {missing}")
            return df
            
        # 重命名列以符合习惯
        df.rename(columns={
            'close': 'Close',
            'high': 'High',
            'low': 'Low',
            'vol': 'Volume'
        }, inplace=True)
        
        # 计算移动平均线
        df['SMA_5'] = df['Close'].rolling(window=5).mean()
        df['SMA_20'] = df['Close'].rolling(window=20).mean()
        
        # 计算相对强弱指数(RSI)
        delta = df['Close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        rs = avg_gain / avg_loss
        df['RSI'] = 100 - (100 / (1 + rs))
        
        # 计算MACD
        df['EMA_12'] = df['Close'].ewm(span=12, adjust=False).mean()
        df['EMA_26'] = df['Close'].ewm(span=26, adjust=False).mean()
        df['MACD'] = df['EMA_12'] - df['EMA_26']
        df['MACD_Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()
        df['MACD_Hist'] = df['MACD'] - df['MACD_Signal']
        
        # 计算成交量变化率
        df['Volume_Change'] = df['Volume'].pct_change()
        
        # 计算价格波动率
        df['Volatility'] = df['Close'].rolling(window=20).std()
        
        # 计算收益率
        df['Return'] = df['Close'].pct_change()
        df['Next_Return'] = df['Return'].shift(-1)
        
        # 标签: 1表示上涨, 0表示下跌
        df['Label'] = np.where(df['Next_Return'] > 0, 1, 0)
        
        return df

    def prepare_features(self, df):
        """准备用于策略的特征"""
        # 填充缺失值
        df.fillna(method='ffill', inplace=True)
        df.fillna(method='bfill', inplace=True)
        
        # 计算财务指标的变化率
        financial_cols = ['ROE', 'Gross_Profit_Margin', 'Net_Profit_Margin', 
                         'Assets_Turnover', 'Debt_To_Assets', 'EPS']
        for col in financial_cols:
            if col in df.columns:
                df[f'{col}_Change'] = df[col].pct_change()
                
        return df

    def price_volume_strategy(self, df):
        """价格成交量策略"""
        signals = pd.Series(index=df.index, data=0)
        
        # 基于均线交叉和RSI的组合信号
        for i in range(20, len(df)):
            # 买入信号: 短期均线上穿长期均线且RSI低于30(超卖)
            if (df['SMA_5'][i] > df['SMA_20'][i] and 
                df['SMA_5'][i-1] <= df['SMA_20'][i-1] and 
                df['RSI'][i] < 30):
                signals[i] = 1
                
            # 卖出信号: 短期均线下穿长期均线且RSI高于70(超买)
            elif (df['SMA_5'][i] < df['SMA_20'][i] and 
                  df['SMA_5'][i-1] >= df['SMA_20'][i-1] and 
                  df['RSI'][i] > 70):
                signals[i] = -1
                
        return signals

    def financial_strategy(self, df):
        """财务指标策略"""
        signals = pd.Series(index=df.index, data=0)
        
        for i in range(40, len(df)):
            # 买入信号: 财务指标良好且价格处于低位
            if ('Net_Profit_Margin' in df.columns and 'ROE' in df.columns and
                df['Net_Profit_Margin'][i] > df['Net_Profit_Margin'].rolling(window=40).mean()[i] and
                df['ROE'][i] > df['ROE'].rolling(window=40).mean()[i] and
                df['RSI'][i] < 40):
                signals[i] = 1
                
            # 卖出信号: 财务指标恶化且价格处于高位
            elif ('Net_Profit_Margin' in df.columns and 'ROE' in df.columns and
                  df['Net_Profit_Margin'][i] < df['Net_Profit_Margin'].rolling(window=40).mean()[i] and
                  df['ROE'][i] < df['ROE'].rolling(window=40).mean()[i] and
                  df['RSI'][i] > 60):
                signals[i] = -1
                
        return signals

    def ml_strategy(self, df):
        """机器学习策略"""
        # 确保有足够的数据用于训练
        if len(df) < 100:
            return pd.Series(index=df.index, data=0)
            
        # 准备特征
        features = ['SMA_5', 'SMA_20', 'RSI', 'MACD', 'MACD_Signal', 'Volume_Change', 'Volatility']
        financial_features = ['ROE', 'Net_Profit_Margin', 'Assets_Turnover', 'Debt_To_Assets']
        for feature in financial_features:
            if feature in df.columns:
                features.append(feature)
                
        df = df.dropna(subset=features + ['Label'])
        if len(df) < 100:
            return pd.Series(index=df.index, data=0)
            
        # 划分训练集和测试集
        X = df[features]
        y = df['Label']
        
        # 使用前80%的数据进行训练，后20%进行测试
        split_index = int(len(df) * 0.8)
        X_train, X_test = X.iloc[:split_index], X.iloc[split_index:]
        y_train, y_test = y.iloc[:split_index], y.iloc[split_index:]
        
        # 训练随机森林分类器
        model = RandomForestClassifier(n_estimators=100, random_state=42)
        model.fit(X_train, y_train)
        
        # 在测试集上预测
        y_pred = model.predict(X_test)
        
        # 生成信号
        signals = pd.Series(index=df.index, data=0)
        signals.iloc[split_index:] = np.where(y_pred == 1, 1, -1)
        
        return signals

    def backtest(self, df, signals, initial_capital=100000):
        """回测策略"""
        # 复制数据并添加信号列
        portfolio = df.copy()
        portfolio['Signal'] = signals
        portfolio['Position'] = portfolio['Signal'].shift(1)
        portfolio['Position'].fillna(0, inplace=True)
        
        # 计算策略收益
        portfolio['Strategy_Return'] = portfolio['Position'] * portfolio['Return']
        
        # 计算累积收益
        portfolio['Cumulative_Return'] = (1 + portfolio['Return']).cumprod()
        portfolio['Cumulative_Strategy_Return'] = (1 + portfolio['Strategy_Return']).cumprod()
        
        # 应用止损和止盈
        portfolio = self.apply_stop_loss_止盈(portfolio)
        
        # 计算最终收益
        final_return = portfolio['Cumulative_Strategy_Return'].iloc[-1] - 1
        
        return portfolio, final_return

    def apply_stop_loss_止盈(self, portfolio, stop_loss=0.08, take_profit=0.15):
        """应用止损和止盈策略"""
        portfolio['Position_Adj'] = portfolio['Position'].copy()
        in_position = False
        entry_price = 0
        
        for i in range(1, len(portfolio)):
            if portfolio['Position'].iloc[i] == 1 and not in_position:
                # 买入信号
                in_position = True
                entry_price = portfolio['Close'].iloc[i]
                portfolio['Position_Adj'].iloc[i] = 1
                
            elif portfolio['Position'].iloc[i] == -1 and in_position:
                # 卖出信号
                in_position = False
                entry_price = 0
                portfolio['Position_Adj'].iloc[i] = -1
                
            elif in_position:
                # 检查止损和止盈
                current_price = portfolio['Close'].iloc[i]
                price_change = (current_price - entry_price) / entry_price
                
                if price_change <= -stop_loss or price_change >= take_profit:
                    # 触发止损或止盈，平仓
                    in_position = False
                    entry_price = 0
                    portfolio['Position_Adj'].iloc[i] = -1
                else:
                    # 保持持仓
                    portfolio['Position_Adj'].iloc[i] = portfolio['Position'].iloc[i-1]
            else:
                # 不在持仓中，保持空仓
                portfolio['Position_Adj'].iloc[i] = 0
                
        # 重新计算策略收益
        portfolio['Strategy_Return_Adj'] = portfolio['Position_Adj'] * portfolio['Return']
        portfolio['Cumulative_Strategy_Return_Adj'] = (1 + portfolio['Strategy_Return_Adj']).cumprod()
        
        return portfolio

    def run(self):
        """运行回测系统"""
        for ticker in self.stock_list:
            print(f"\nProcessing {ticker}...")
            df = self.fetch_data(ticker)
            if df is None or df.empty:
                continue
                
            df = self.calculate_technical_indicators(df)
            df = self.prepare_features(df)
            
            self.results[ticker] = {}
            
            for strategy_name, strategy_func in self.strategies.items():
                signals = strategy_func(df)
                portfolio, final_return = self.backtest(df, signals)
                
                self.results[ticker][strategy_name] = {
                    'portfolio': portfolio,
                    'final_return': final_return,
                    'signals': signals
                }
                
                print(f"{strategy_name} strategy for {ticker} final return: {final_return:.2%}")

    def plot_results(self, ticker):
        """绘制回测结果"""
        if ticker not in self.results:
            print(f"No results found for {ticker}")
            return
            
        plt.figure(figsize=(14, 10))
        
        for i, (strategy_name, result) in enumerate(self.results[ticker].items(), 1):
            plt.subplot(len(self.strategies), 1, i)
            portfolio = result['portfolio']
            
            plt.plot(portfolio.index, portfolio['Cumulative_Return'], label='Buy & Hold')
            plt.plot(portfolio.index, portfolio['Cumulative_Strategy_Return'], label=f'{strategy_name} Strategy')
            plt.plot(portfolio.index, portfolio['Cumulative_Strategy_Return_Adj'], label=f'{strategy_name} Strategy (with SL/TP)')
            
            plt.title(f'{ticker} Cumulative Returns - {strategy_name} Strategy')
            plt.xlabel('Date')
            plt.ylabel('Cumulative Return')
            plt.legend()
            plt.grid(True)
            
        plt.tight_layout()
        plt.show()

    def summary(self):
        """生成回测结果摘要"""
        summary_data = []
        
        for ticker, strategies in self.results.items():
            for strategy_name, result in strategies.items():
                summary_data.append({
                    'Ticker': ticker,
                    'Strategy': strategy_name,
                    'Return': result['final_return']
                })
                
        summary_df = pd.DataFrame(summary_data)
        return summary_df

# 运行示例
if __name__ == "__main__":
    # 请替换为你的Tushare token
    token = "1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c"
    
    # 选择一些股票进行回测
    # 注意：Tushare使用的股票代码格式为"股票代码+交易所代码"，例如"000001.SZ"
    stock_list = ['000001.SZ', '600519.SH', '601318.SH', '002594.SZ', '300750.SZ']
    start_date = '20200101'
    end_date = '20231231'
    
    quant_system = QuantSystem(stock_list, start_date, end_date, token)
    quant_system.run()
    
    # 显示结果摘要
    summary = quant_system.summary()
    print("\n策略回测结果汇总:")
    print(summary.sort_values(by=['Ticker', 'Return'], ascending=[True, False]))
    
    # 绘制第一支股票的结果图表
    if stock_list:
        quant_system.plot_results(stock_list[0])


# In[ ]:




